import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm

from utils import search_data
import os
from datetime import datetime

gongye_data = pd.read_excel(f"cache/{search_data('生产测_工业')}", 
    index_col=0)['工业增加值当月同比（合并1-2月）']

fuwuye_data = pd.read_excel(f"cache/{search_data('生产测_服务业')}", 
    index_col=0)['服务业生产指数：当月同比']

xf_data = pd.read_excel(f"cache/{search_data('需求侧_消费')}", 
    index_col=0)['社会消费品零售总额：当月同比']

tz_data = pd.read_excel(f"cache/{search_data('需求侧_投资')}", 
    index_col=0)['中国:固定资产投资完成额:累计同比']

chukou_data = pd.read_excel(f"cache/{search_data('需求侧_进出口')}", 
    index_col=0)['出口金额:当月同比']

df_plot = pd.concat([gongye_data, fuwuye_data, xf_data, tz_data, chukou_data], axis=1)
df_plot.columns = [
    '工业增加值当月同比', '服务业生产指数当月同比', 
    '社会消费品零售总额当月同比', '固定资产投资完成额累计同比',
    '出口金额当月同比'
    ]
df_plot.index = pd.to_datetime(df_plot.index)

df_cum = pd.concat(
    [
        pd.read_excel(f"cache/{search_data('生产测_工业')}", index_col=0)['工业增加值累计同比'],
        pd.read_excel(f"cache/{search_data('生产测_服务业')}", index_col=0)['服务业生产指数：累计同比'],
        pd.read_excel(f"cache/{search_data('需求侧_消费')}", index_col=0)['社会消费品零售总额：累计同比'],
        pd.read_excel(f"cache/{search_data('需求侧_投资')}", index_col=0)['中国:固定资产投资完成额:累计同比'],
        pd.read_excel(f"cache/{search_data('需求侧_进出口')}", index_col=0)['出口金额:累计同比']
    ], axis=1
)
df_cum.columns = [
    '工业增加值累计同比', '服务业生产指数累计同比', 
    '社会消费品零售总额累计同比', '固定资产投资完成额累计同比',
    '出口金额累计同比'
    ]
df_cum.index = pd.to_datetime(df_cum.index)

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def generate_color_gradient(base_color, n):
    """
    根据主颜色和数量生成从浅到深的渐变色列表
    
    参数:
    base_color: str - 主颜色，可以是中文颜色名称如'红色'，或英文颜色名称如'red'，或十六进制颜色码如'#ff0000'
    n: int - 需要生成的颜色数量
    
    返回:
    list - 包含n个十六进制颜色码的列表，从浅到深
    """
    # 颜色名称映射（中英文）
    color_mapping = {
        'red': '#ff0000', '红色': '#ff0000',
        'blue': '#1f77b4', '蓝色': '#1f77b4',
        'green': '#2ca02c', '绿色': '#2ca02c',
        'orange': '#ff7f0e', '橙色': '#ff7f0e',
        'purple': '#9467bd', '紫色': '#9467bd',
        'brown': '#8c564b', '棕色': '#8c564b',
        'pink': '#e377c2', '粉色': '#e377c2',
        'gray': '#7f7f7f', '灰色': '#7f7f7f',
        'yellow': '#bcbd22', '黄色': '#bcbd22',
        'cyan': '#17becf', '青色': '#17becf'
    }
    
    # 将颜色名称转换为十六进制
    if base_color.lower() in color_mapping:
        hex_color = color_mapping[base_color.lower()]
    elif base_color in color_mapping:
        hex_color = color_mapping[base_color]
    elif base_color.startswith('#') and len(base_color) in [7, 9]:
        hex_color = base_color
    else:
        # 默认红色
        hex_color = '#ff0000'
    
    # 十六进制转RGB
    hex_color = hex_color.lstrip('#')
    if len(hex_color) == 6:
        r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    else:
        # 处理RGBA
        r, g, b = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
    
    # 生成渐变色 - 从浅到深
    colors = []
    for i in range(n):
        # 计算当前颜色的亮度因子，从0.3（浅）到1.0（深）
        factor = 0.3 + (i / (n - 1)) * 0.7 if n > 1 else 0.65
        
        # 根据因子调整RGB值
        # 对于浅色，混合更多白色；对于深色，更接近原始颜色
        light_r = int(255 - (255 - r) * factor)
        light_g = int(255 - (255 - g) * factor)
        light_b = int(255 - (255 - b) * factor)
        
        # 转换回十六进制
        light_hex = '#{:02x}{:02x}{:02x}'.format(light_r, light_g, light_b)
        colors.append(light_hex)
    
    return colors


def generate_multiple_color_gradients(column_names, n_per_column):
    """
    为多列数据生成不同的色系，确保各列颜色有明显区分
    
    参数:
    column_names: list - 列名列表
    n_per_column: int - 每列需要生成的颜色数量
    
    返回:
    dict - 键为列名，值为该列的颜色列表
    """
    # 预定义的主色系，按照视觉区分度排序
    # 第一组：高对比度基础色
    primary_colors = [
        'blue', 'red', 'green', 'orange', 'purple', 
        'brown', 'pink', 'cyan', 'yellow', 'gray'
    ]
    
    # 第二组：补充色，用于更多列的情况
    secondary_colors = [
        '#ff6b6b', '#4ecdc4', '#45b7d1', '#96ceb4', '#ffeaa7',
        '#dfe6e9', '#fdcb6e', '#fd79a8', '#6c5ce7', '#a29bfe'
    ]
    
    # 合并所有可用颜色
    all_colors = primary_colors + secondary_colors
    
    # 根据列名的重要性或特定关键词优先分配更突出的颜色
    # 定义关键词映射到特定颜色的优先级
    keyword_color_mapping = {
        '工业': 'blue',
        '服务': 'red',
        '生产': 'green',
        '消费': 'orange',
        '投资': 'purple',
        '出口': 'cyan',
        '进口': 'brown',
        '金融': 'pink',
        '价格': 'yellow',
        'CPI': 'red',
        'PPI': 'orange'
    }
    
    # 为每列分配一个主色系
    color_gradients = {}
    used_colors = set()
    
    # 首先处理有特定关键词的列名
    for column in column_names:
        # 查找列名中是否包含关键词
        assigned = False
        for keyword, preferred_color in keyword_color_mapping.items():
            if keyword in column and preferred_color not in used_colors:
                color_gradients[column] = generate_color_gradient(preferred_color, n_per_column)
                used_colors.add(preferred_color)
                assigned = True
                break
        
        # 如果没有匹配的关键词，则稍后处理
        if not assigned:
            color_gradients[column] = None
    
    # 处理剩余的列名
    for column in column_names:
        if color_gradients[column] is None:
            # 找到下一个未使用的颜色
            for i in range(len(all_colors)):
                # 使用索引取模确保不会超出颜色列表范围
                color_index = i % len(all_colors)
                color = all_colors[color_index]
                
                # 确保颜色未被使用
                if isinstance(color, str):
                    # 颜色名称的情况
                    if color not in used_colors:
                        color_gradients[column] = generate_color_gradient(color, n_per_column)
                        used_colors.add(color)
                        break
                elif color.startswith('#'):
                    # 十六进制颜色码的情况
                    if color not in used_colors:
                        color_gradients[column] = generate_color_gradient(color, n_per_column)
                        used_colors.add(color)
                        break
    
    return color_gradients



# 绘制多指标柱状图
def plot_multi_bar(df, title='多指标柱状图', figsize=(14, 8), save_path=None, 
                  colors=None, show_values=False, add_grid=True, transpose=False):
    """
    绘制美化的多指标柱状图
    
    参数:
    df: DataFrame, 包含多个指标的数据
    title: str, 图表标题
    figsize: tuple, 图表大小
    save_path: str, 保存路径，None表示自动生成
    colors: list, 柱状图颜色列表，None表示使用默认颜色
    show_values: bool, 是否显示数值标签
    add_grid: bool, 是否添加网格线
    transpose: bool, 是否转置绘图方向（True表示以指标为x轴，日期为系列）
    """
    # 创建图形
    fig, ax = plt.subplots(figsize=figsize)
    
    # 根据transpose参数决定数据结构
    if transpose:
        # 转置数据：指标为x轴，日期为系列
        plot_df = df.copy()
        
        # 确保索引是日期类型并排序
        if not pd.api.types.is_datetime64_any_dtype(plot_df.index):
            try:
                plot_df.index = pd.to_datetime(plot_df.index)
            except:
                # 如果无法转换为日期，保持原样
                pass
        
        # 按日期排序（如果是日期索引）
        if pd.api.types.is_datetime64_any_dtype(plot_df.index):
            plot_df = plot_df.sort_index()
        
        categories = plot_df.columns  # 指标作为类别（x轴）
        series = plot_df.index  # 日期作为系列（图例）
        n_series = len(series)
        width = 0.8 / n_series
        
        # 为每列指标生成不同的色系
        if colors is None and len(categories) > 1:
            # 使用新的多色系生成函数
            color_gradients = generate_multiple_color_gradients(categories, n_series)
        else:
            # 默认颜色列表（更多颜色选项）
            default_colors = [
                '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
                '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf',
                '#aec7e8', '#ffbb78', '#98df8a', '#ff9896', '#c5b0d5',
                '#c49c94', '#f7b6d2', '#c7c7c7', '#dbdb8d', '#9edae5'
            ]
            # 循环使用颜色，确保有足够的颜色
            colors = [default_colors[i % len(default_colors)] for i in range(n_series)]
        
        # 绘制柱状图
        for i, date in enumerate(series):
            # 将日期格式化为字符串作为图例标签
            if pd.api.types.is_datetime64_any_dtype(type(date)):
                date_label = date.strftime('%Y-%m')
            elif hasattr(date, 'strftime'):
                date_label = date.strftime('%Y-%m')
            else:
                date_label = str(date)
            
            # 为每个指标类别绘制对应的柱子
            for j, category in enumerate(categories):
                x = j + i * width
                
                # 确保数据有效
                try:
                    value = plot_df.loc[date, category]
                    
                    # 选择颜色
                    if colors is not None:
                        # 使用传入的颜色列表
                        color = colors[i % len(colors)]
                    elif len(categories) > 1:
                        # 使用为该类别生成的色系
                        color = color_gradients[category][i]
                    else:
                        # 单类别情况使用默认颜色
                        color = '#1f77b4'
                    
                    # 只有第一个类别添加图例标签，避免重复
                    label = date_label if j == 0 else ""
                    
                    bar = ax.bar(x, value, width, label=label, color=color, alpha=0.85)
                    
                    # 显示数值标签
                    if show_values:
                        height = bar[0].get_height()
                        ax.annotate(f'{height:.1f}',
                                    xy=(x, height),
                                    xytext=(0, 3),  # 3点垂直偏移
                                    textcoords="offset points",
                                    ha='center', va='bottom', fontsize=9)
                    
                except (KeyError, ValueError):
                    # 处理可能的索引错误或值错误
                    continue
            
            # 显示数值标签
            if show_values:
                for rect in bar:
                    height = rect.get_height()
                    ax.annotate(f'{height:.1f}',
                                xy=(rect.get_x() + rect.get_width() / 2, height),
                                xytext=(0, 3),  # 3点垂直偏移
                                textcoords="offset points",
                                ha='center', va='bottom', fontsize=9)
        
        # 设置x轴标签
        # 调整x轴位置，使柱状图居中
        ax.set_xticks([j + width * (n_series - 1) / 2 for j in range(len(categories))])
        
        # 优化x轴标签显示
        ax.set_xticklabels(categories, fontsize=10, rotation=0, ha='center')
        
        # 根据指标名称长度调整标签显示
        if any(len(cat) > 10 for cat in categories):
            ax.tick_params(axis='x', labelsize=9)
            
        # 设置更适合的x轴标签
        ax.set_xlabel('指标类别', fontsize=12, labelpad=10)
        ax.set_ylabel('同比增长率(%)', fontsize=12, labelpad=10)
        
        # 设置标题
        if title:
            ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
        else:
            # 自动生成标题
            ax.set_title('各指标在不同日期的对比', fontsize=16, fontweight='bold', pad=20)
    else:
        n_series = len(df.columns)
        # 原有逻辑：日期为x轴，指标为系列
        # 设置柱状图宽度
        width = 0.8 / len(df.columns)
        
        # 获取数据
        index = df.index
        columns = df.columns
        
        # 默认颜色列表
        default_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
        if colors is None:
            colors = default_colors[:len(columns)]
        
        # 绘制柱状图
        bars = []
        for i, (column, color) in enumerate(zip(columns, colors)):
            x = [j + i * width for j in range(len(index))]
            bar = ax.bar(x, df[column], width, label=column, color=color, alpha=0.85)
            bars.append(bar)
            
            # 显示数值标签
            if show_values:
                for rect in bar:
                    height = rect.get_height()
                    ax.annotate(f'{height:.1f}',
                                xy=(rect.get_x() + rect.get_width() / 2, height),
                                xytext=(0, 3),  # 3点垂直偏移
                                textcoords="offset points",
                                ha='center', va='bottom', fontsize=9)
        
        # 设置x轴标签
        ax.set_xticks([j + width * (len(columns) - 1) / 2 for j in range(len(index))])
        ax.set_xticklabels([date.strftime('%Y-%m') for date in index], rotation=45, ha='right')
        
        # 添加标题和标签（美化样式）
        ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
        ax.set_xlabel('日期', fontsize=12, labelpad=10)
        ax.set_ylabel('同比增长率(%)', fontsize=12, labelpad=10)
    
    # 添加图例
    if n_series <= 2:  # 较少系列时使用右侧位置
        ax.legend(loc='upper right', fontsize=10, frameon=False)
    else:
        # 图例过多时，使用右侧垂直排列，确保与图像无冲突
        # 使用bbox_to_anchor将图例放在图的右侧外部
        ax.legend(loc='center left', bbox_to_anchor=(1.02, 0.5), 
                  ncol=1, fontsize=8, frameon=False, 
                  title='图例', title_fontsize=9)
        # 调整subplot布局，为右侧图例留出空间
        plt.tight_layout(rect=[0, 0, 0.85, 1])  # 左侧、底部、右侧、顶部的留白比例
    
    # 添加网格线
    if add_grid:
        ax.grid(axis='y', linestyle='--', alpha=0.7)
    
    # 美化图表边框
    for spine in ['top', 'right']:
        ax.spines[spine].set_visible(False)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图表
    if save_path:
        # 确保保存目录存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"图表已保存至: {save_path}")
    else:
        # 自动生成保存路径
        today = datetime.now().strftime('%Y-%m-%d')
        auto_dir = f"../plot/plot_{today}"
        os.makedirs(auto_dir, exist_ok=True)
        auto_save_path = f"{auto_dir}/{title}_t.png" if transpose else f"{auto_dir}/{title}.png"
        plt.savefig(auto_save_path, dpi=300, bbox_inches='tight')
        print(f"图表已自动保存至: {auto_save_path}")
    
    # 显示图表
    plt.show()
    plt.close()

# 使用函数绘制df_supply的多指标柱状图（转置方式：指标为x轴，日期为系列）
# 测试多色系功能，不传入colors参数，让系统自动为每列生成独立色系
plot_multi_bar(
    df_plot.iloc[-8:-1,:],  # 使用最近12条数据
    title='生产测（工业服务业）与需求侧（消费投资出口）月度数据对比',
    show_values=True,  # 显示数值标签
    add_grid=True,  # 添加网格线
    transpose=True  # 设置为True，以指标为x轴，日期为系列
)
plot_multi_bar(
    df_cum.iloc[-8:-1,:],  # 使用最近12条数据
    title='生产测（工业服务业）与需求侧（消费投资出口）月度数据对比(累计同比)',
    show_values=True,  # 显示数值标签
    add_grid=True,  # 添加网格线
    transpose=True  # 设置为True，以指标为x轴，日期为系列
)
# 测试原始绘图方式（日期为x轴）
# plot_multi_bar(
#     df_plot.iloc[-8:,:],  # 使用最近12条数据
#     title='对比',
#     show_values=True,  # 显示数值标签
#     add_grid=True,  # 添加网格线
#     transpose=False  # 以日期为x轴
# )